import os
import torch
import csv
import time
import numpy as np
import mindspore as ms
from src import utils
from src.pytorch_init import PytorchInit
from src.mindspore_init import MindsporeInit
from src.compare_algorithm import alg_similarity
from src.compare_algorithm import alg_cosine_similarity
from src.compare_algorithm import alg_variance
from src.compare_algorithm import alg_max_absolute_error
from src.compare_algorithm import alg_max_relative_error
from src.compare_algorithm import alg_mean_absolute_error
from src.compare_algorithm import alg_mean_relative_error
from src.compare_algorithm import alg_mean_square_error
from src.compare_algorithm import alg_relative_euclidean_distance


MS_OUTPUT_PATH = "data/output/MindSpore"
PT_OUTPUT_PATH = "data/output/PyTorch"
INPUT_DATA_PATH = "data/input"
INPUT_PATH = "data/input"
RESULT_COLUMNS = ["Pytorch data", "Pytorch output shape", "MindSpore data", "MindSpore output shape",
                  "cosine similarity", "similarity", "variance", "max absolute error", "max relative error",
                  "mean absolute error", "mean relative error", "mean square error", "relative euclidean distance"]


class Compare:
    def __init__(self, ptmodule, msmodule, input_shape, input_dtype, init_mode, out_path, print_result):
        self.ptmodule = ptmodule
        self.msmodule = msmodule
        self.input_shape = input_shape
        self.input_dtype = input_dtype
        self.init_mode = init_mode
        self.print_result = print_result
        self.pt_init = PytorchInit()
        self.ms_init = MindsporeInit()
        self.out_path = out_path
        self.pt_out_path = "NAN"
        self.ms_out_path = "NAN"
        self.pt_out_shape = "NAN"
        self.ms_out_shape = "NAN"
        
    # generate input_data for net
    def gen_input_data(self):
        utils.print_info_log("==============================Start generate input data==============================")
        input_data = []
        index = 1
        for shape, dtype in zip(self.input_shape, self.input_dtype):
            if dtype == "int32":
                input_data_item = np.random.randint(1, 10, shape).astype(dtype)
            else:
                input_data_item = np.random.random_sample(shape).astype(dtype)
            input_data.append(input_data_item)
            input_name = "input_%s.npy" % index
            input_data_path = os.path.join(self.out_path, INPUT_DATA_PATH, input_name)
            utils.save_numpy_data(input_data_path, input_data_item)
            utils.print_info_log("Generate input data at: %s" % input_data_path)
            index += 1
        if len(self.input_shape) != len(input_data):
            utils.print_error_log("input_shape count not match generated input data.")
            raise utils.CompareException(utils.ACCURACY_COMPARISON_INVALID_PARAM_ERROR)
        return input_data

    def run_pt_net(self, module, input_data_list):
        param_list = []
        for input_data, dtype in zip(input_data_list, self.input_dtype):
            x = utils.TORCH_MAP.get(dtype)(input_data)
            param_list.append(x)
        result = module(*param_list)
        return result

    def run_ms_net(self, module, input_data_list):
        param_list = []
        for input_data, dtype in zip(input_data_list, self.input_dtype):
            x = ms.Tensor(input_data, utils.MS_MAP.get(dtype))
            param_list.append(x)
        result = module(*param_list)
        return result

    # generate init constant for both net
    def gen_init_constant(self):
        constant = 1
        if self.init_mode == "random":
            constant = np.random.rand()
        utils.print_info_log("Init net condtant is: %s" % constant)
        return constant
    
    def compare_results(self, result_pt_list, result_ms_list):
        index = 0
        compare_result = []
        for (result_pt, result_ms) in zip(result_pt_list, result_ms_list):
            result_pt = result_pt.detach().numpy()
            result_ms = result_ms.asnumpy()
            self.check_out_data_shape(index, result_ms, result_pt)
            if self.pt_out_shape != self.ms_out_shape:
                utils.print_error_log("output results are in different shapes!")
                raise utils.CompareException(utils.ACCURACY_COMPARISON_NOT_SUPPORT_ERROR)
            else:
                self.save_out_data(index, result_ms, result_pt)
                result_pt = result_pt.reshape(1, -1)
                result_ms = result_ms.reshape(1, -1)

                result = self.compare_data(result_ms, result_pt, index)
                compare_result.append(result)
            index += 1
        return compare_result

    def save_out_data(self, index, result_ms, result_pt):
        utils.print_info_log("================= ======Start Save Out Data %s =========================" % index)
        result_file = "output_%s.npy" % index
        self.ms_out_path = os.path.join(self.out_path, MS_OUTPUT_PATH, result_file)
        utils.save_numpy_data(self.ms_out_path, result_ms)
        utils.print_info_log("Saved MindSpore output data at: %s" % self.ms_out_path)
        self.pt_out_path = os.path.join(self.out_path, PT_OUTPUT_PATH, result_file)
        utils.save_numpy_data(self.pt_out_path, result_pt)
        utils.print_info_log("Saved PyTorch output data at: %s" % self.pt_out_path)

    def check_out_data_shape(self, index, result_ms, result_pt):
        utils.print_info_log("=========================Start Check Out Data %s ===========================" % index)
        self.pt_out_shape = result_pt.shape
        self.ms_out_shape = result_pt.shape
        utils.print_info_log("shape of result_pt: %s" % str(self.pt_out_shape))
        utils.print_info_log("shape of result_ms: %s" % str(self.ms_out_shape))
        if self.print_result:
            print("-result_pt-: \n", result_pt)
            print("-result_ms-: \n", result_ms)

    def compare_data(self, result_ms, result_pt, index):
        utils.print_info_log("=========================Start Compare Out Data %s ===========================" % index)
        similarity = alg_similarity(result_pt, result_ms)
        cos_sim = alg_cosine_similarity(result_pt, result_ms)
        variance = alg_variance(result_pt, result_ms)
        max_absolute_error = alg_max_absolute_error(result_pt, result_ms)
        max_relative_error = alg_max_relative_error(result_pt, result_ms)
        mean_absolute_error = alg_mean_absolute_error(result_pt, result_ms)
        mean_relative_error = alg_mean_relative_error(result_pt, result_ms)
        mean_square_error = alg_mean_square_error(result_pt, result_ms)
        relative_euclidean_distance = alg_relative_euclidean_distance(result_pt, result_ms)
        print("-similarity-: ", similarity)
        print("-cosine similarity-: ", cos_sim)
        print("-variance-: ", variance)
        print("-max_absolute_error-", max_absolute_error)
        print("-max_relative_error-", max_relative_error)
        print("-mean_absolute_error-", mean_absolute_error)
        print("-mean_relative_error-", mean_relative_error)
        print("-mean_square_error-", mean_square_error)
        print("-relative_euclidean_distance-", relative_euclidean_distance)

        result = [self.pt_out_path, self.pt_out_shape, self.ms_out_path, self.ms_out_shape, cos_sim, similarity,
                  variance, max_absolute_error, max_relative_error, mean_absolute_error, mean_relative_error,
                  mean_square_error, relative_euclidean_distance]
        return result

    def start_compare(self):
        pt_net, ms_net = self.init_net()
        self.check_input_shape()
        input_data = self.gen_input_data()
        result_ms, result_pt = self.inference_net(input_data, ms_net, pt_net)
        self.check_output(result_ms, result_pt)
        compare_results = self.compare_results(result_pt, result_ms)
        self.save_result(compare_results)

    def save_result(self, compare_results):
        utils.print_info_log("=================================Start save result=================================")
        result_path = os.path.join(self.out_path, "compare_result.csv")
        with open(result_path, "a+") as csvfile:
            writer = csv.writer(csvfile)

            # 先写入columns_name
            writer.writerow(RESULT_COLUMNS)
            # 写入多行用writerows
            writer.writerows(compare_results)
            utils.print_info_log("The comparison result have been written to %s" % result_path)

    @staticmethod
    def check_output(result_ms, result_pt):
        ms_result_num = len(result_ms)
        utils.print_info_log("The MindSpore net inference have %s result." % ms_result_num)
        pt_result_num = len(result_pt)
        utils.print_info_log("The PyTorch net inference have %s result." % pt_result_num)
        if ms_result_num != pt_result_num:
            utils.print_error_log("output results are in different counts!")
            raise utils.CompareException(utils.ACCURACY_COMPARISON_NOT_SUPPORT_ERROR)

    def inference_net(self, input_data, ms_net, pt_net):
        utils.print_info_log("=================================Start inference net=================================")
        start_pt = time.time()
        result_pt = self.run_pt_net(pt_net, input_data)
        end_pt = time.time()
        utils.print_info_log("The PyTorch net inference completed cost %s seconds." % (end_pt - start_pt))
        start_ms = time.time()
        result_ms = self.run_ms_net(ms_net, input_data)
        end_ms = time.time()
        utils.print_info_log("The MindSpore net inference completed cost %s seconds." % (end_ms - start_ms))
        return result_ms, result_pt

    def init_net(self):
        utils.print_info_log("=================================Start init Net=================================")
        constant = self.gen_init_constant()
        pt_module = self.pt_init.net_init(self.ptmodule, constant)
        ms_module = self.ms_init.net_init(self.msmodule, constant)
        return pt_module, ms_module

    def check_input_shape(self):
        utils.print_info_log("==========================Start check and parse input shape===========================")
        if not isinstance(self.input_shape, list):
            utils.print_error_log("input_shape should be a list.")
            raise utils.CompareException(utils.ACCURACY_COMPARISON_INVALID_PARAM_ERROR)
        for item in self.input_shape:
            if not isinstance(item, list):
                utils.print_error_log("each input_shape should set be in a list, Separate multiple "
                                      "inputs with commas(,) in list. E.g: one input [[1,2]]; "
                                      "two inputs [[1,2],[1,2]]")
                raise utils.CompareException(utils.ACCURACY_COMPARISON_INVALID_PARAM_ERROR)
            for dim in item:
                if not isinstance(dim, int):
                    utils.print_error_log("The value(%s) of input_shape not int." % item)
                    raise utils.CompareException(utils.ACCURACY_COMPARISON_INVALID_PARAM_ERROR)
            utils.print_info_log("one of input data shape is %s" % item)
        utils.print_info_log("After check input_shape, this compare task have %s input" % len(self.input_shape))
